from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
import asyncio
import json
import random
from typing import List, Dict, Any
import uvicorn

app = FastAPI()

# 允许跨域
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class WarUnit:
    def __init__(self, unit_type: str, strength: int, morale: int, side: str, row: int, col: int):
        self.type = unit_type
        self.strength = strength
        self.morale = morale
        self.side = side
        self.row = row
        self.col = col
        self.max_strength = strength
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "type": self.type,
            "strength": self.strength,
            "morale": self.morale,
            "side": self.side,
            "row": self.row,
            "col": self.col,
            "max_strength": self.max_strength
        }

class WarSimulator:
    def __init__(self, grid_size: int = 10):
        self.grid_size = grid_size
        self.grid = [[None for _ in range(grid_size)] for _ in range(grid_size)]
        self.red_army: List[WarUnit] = []
        self.blue_army: List[WarUnit] = []
        self.current_turn = 0
        self.is_running = False
        self.max_turns = 50
        
    def reset(self):
        """重置战场"""
        self.grid = [[None for _ in range(self.grid_size)] for _ in range(self.grid_size)]
        self.red_army.clear()
        self.blue_army.clear()
        self.current_turn = 0
        self.is_running = False
        
    def deploy_armies(self):
        """部署初始军队"""
        # 红方军队（左侧）
        for i in range(5):
            unit = WarUnit(
                unit_type="步兵",
                strength=20,
                morale=100,
                side="red",
                row=random.randint(0, self.grid_size - 1),
                col=random.randint(0, 2)
            )
            self.red_army.append(unit)
            self.grid[unit.row][unit.col] = unit
            
        # 蓝方军队（右侧）
        for i in range(5):
            unit = WarUnit(
                unit_type="步兵",
                strength=20,
                morale=100,
                side="blue",
                row=random.randint(0, self.grid_size - 1),
                col=self.grid_size - 1 - random.randint(0, 2)
            )
            self.blue_army.append(unit)
            self.grid[unit.row][unit.col] = unit
    
    def recruit_militia(self) -> List[Dict[str, Any]]:
        """招募民兵"""
        events = []
        
        # 红方招募
        if random.random() < 0.3:
            row = random.randint(0, self.grid_size - 1)
            col = 0
            if self.grid[row][col] is None:
                unit = WarUnit(
                    unit_type="民兵",
                    strength=10,
                    morale=80,
                    side="red",
                    row=row,
                    col=col
                )
                self.red_army.append(unit)
                self.grid[row][col] = unit
                events.append({
                    "type": "recruit",
                    "side": "red",
                    "unit": unit.to_dict()
                })
        
        # 蓝方招募
        if random.random() < 0.3:
            row = random.randint(0, self.grid_size - 1)
            col = self.grid_size - 1
            if self.grid[row][col] is None:
                unit = WarUnit(
                    unit_type="民兵",
                    strength=10,
                    morale=80,
                    side="blue",
                    row=row,
                    col=col
                )
                self.blue_army.append(unit)
                self.grid[row][col] = unit
                events.append({
                    "type": "recruit",
                    "side": "blue",
                    "unit": unit.to_dict()
                })
        
        return events
    
    def move_units(self) -> List[Dict[str, Any]]:
        """移动单位"""
        events = []
        
        # 红方移动（向右）
        for unit in self.red_army[:]:
            if unit.strength <= 0:
                continue
                
            new_col = min(unit.col + 1, self.grid_size - 1)
            if new_col != unit.col and self.grid[unit.row][new_col] is None:
                self.grid[unit.row][unit.col] = None
                unit.col = new_col
                self.grid[unit.row][unit.col] = unit
                events.append({
                    "type": "move",
                    "unit": unit.to_dict()
                })
        
        # 蓝方移动（向左）
        for unit in self.blue_army[:]:
            if unit.strength <= 0:
                continue
                
            new_col = max(unit.col - 1, 0)
            if new_col != unit.col and self.grid[unit.row][new_col] is None:
                self.grid[unit.row][unit.col] = None
                unit.col = new_col
                self.grid[unit.row][new_col] = unit
                events.append({
                    "type": "move",
                    "unit": unit.to_dict()
                })
        
        return events
    
    def combat_phase(self) -> List[Dict[str, Any]]:
        """战斗阶段"""
        events = []
        
        # 检查所有相邻的敌对单位
        for row in range(self.grid_size):
            for col in range(self.grid_size):
                unit = self.grid[row][col]
                if unit is None or unit.strength <= 0:
                    continue
                
                # 寻找相邻的敌人
                enemies = self.find_adjacent_enemies(row, col, unit.side)
                if enemies:
                    # 随机选择一个敌人战斗
                    enemy = random.choice(enemies)
                    
                    # 计算伤害
                    damage_to_enemy = random.randint(1, int(unit.strength * 0.5))
                    damage_to_self = random.randint(1, int(enemy.strength * 0.5))
                    
                    # 应用伤害
                    enemy.strength = max(0, enemy.strength - damage_to_enemy)
                    unit.strength = max(0, unit.strength - damage_to_self)
                    
                    events.append({
                        "type": "combat",
                        "attacker": unit.to_dict(),
                        "defender": enemy.to_dict(),
                        "damage_to_defender": damage_to_enemy,
                        "damage_to_attacker": damage_to_self
                    })
                    
                    # 检查单位是否被消灭
                    if enemy.strength <= 0:
                        self.grid[enemy.row][enemy.col] = None
                        events.append({
                            "type": "eliminated",
                            "unit": enemy.to_dict()
                        })
                    
                    if unit.strength <= 0:
                        self.grid[unit.row][unit.col] = None
                        events.append({
                            "type": "eliminated",
                            "unit": unit.to_dict()
                        })
        
        return events
    
    def find_adjacent_enemies(self, row: int, col: int, side: str) -> List[WarUnit]:
        """寻找相邻的敌人"""
        enemies = []
        directions = [(-1, 0), (1, 0), (0, -1), (0, 1)]
        
        for dr, dc in directions:
            new_row, new_col = row + dr, col + dc
            if (0 <= new_row < self.grid_size and 
                0 <= new_col < self.grid_size):
                enemy = self.grid[new_row][new_col]
                if enemy and enemy.side != side and enemy.strength > 0:
                    enemies.append(enemy)
        
        return enemies
    
    def check_victory(self) -> Dict[str, Any]:
        """检查胜利条件"""
        red_alive = [u for u in self.red_army if u.strength > 0]
        blue_alive = [u for u in self.blue_army if u.strength > 0]
        
        if len(red_alive) == 0:
            return {"winner": "blue", "reason": "红方全军覆没"}
        elif len(blue_alive) == 0:
            return {"winner": "red", "reason": "蓝方全军覆没"}
        elif self.current_turn >= self.max_turns:
            # 比较剩余兵力
            red_total = sum(u.strength for u in red_alive)
            blue_total = sum(u.strength for u in blue_alive)
            
            if red_total > blue_total:
                return {"winner": "red", "reason": "兵力优势"}
            elif blue_total > red_total:
                return {"winner": "blue", "reason": "兵力优势"}
            else:
                return {"winner": "draw", "reason": "平局"}
        
        return None
    
    def get_game_state(self) -> Dict[str, Any]:
        """获取游戏状态"""
        return {
            "grid_size": self.grid_size,
            "current_turn": self.current_turn,
            "red_units": [u.to_dict() for u in self.red_army if u.strength > 0],
            "blue_units": [u.to_dict() for u in self.blue_army if u.strength > 0],
            "red_total_strength": sum(u.strength for u in self.red_army),
            "blue_total_strength": sum(u.strength for u in self.blue_army),
            "is_running": self.is_running
        }
    
    async def execute_turn(self) -> Dict[str, Any]:
        """执行一个回合"""
        self.current_turn += 1
        
        # 执行各个阶段
        recruit_events = self.recruit_militia()
        move_events = self.move_units()
        combat_events = self.combat_phase()
        
        # 检查胜利
        victory = self.check_victory()
        
        return {
            "turn": self.current_turn,
            "recruit_events": recruit_events,
            "move_events": move_events,
            "combat_events": combat_events,
            "victory": victory,
            "game_state": self.get_game_state()
        }

# 全局模拟器实例
simulator = WarSimulator()

class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []
    
    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)
    
    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)
    
    async def broadcast(self, message: Dict[str, Any]):
        for connection in self.active_connections[:]:
            try:
                await connection.send_json(message)
            except:
                self.active_connections.remove(connection)

manager = ConnectionManager()

@app.websocket("/ws")
async def websocket